import shutil
from typing import List

from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import true
from sqlalchemy.orm import Session
from vsearcher import VSearcher
from app.models.base import BaseModel


class Mock:

    def __init__(self, db: SQLAlchemy, reset=True):
        self.db = db
        self.reset = reset

    def get_session(self) -> Session:
        return self.db.session

    def create(self, refer_class: BaseModel, num: int, data_dict: dict) -> List[BaseModel]:
        ''' 返回{num}个用{data_dict}初始化的{refer_class}对象列表 '''

        def temp(data_dict: dict, i):
            return refer_class().set_attr({k: f'{v}{i}' if isinstance(v, str) else v for k, v in data_dict.items()})
        return [temp(data_dict, i + 1) for i in range(num)]

    def __enter__(self):
        if self.reset:
            self.db.drop_all()
            self.db.create_all()
            shutil.rmtree(VSearcher.config.RootPath.output_dir,
                          ignore_errors=true)
            shutil.rmtree(
                VSearcher.config.RootPath.output_search_result_dir, ignore_errors=True)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.db.session.commit()
        return False
